#include "yolov6_detector.h"

namespace {

constexpr int kNumClasses = 80;
constexpr float kNmsThresh = 0.45;
constexpr float kObjectnessThresh = 0.45;
constexpr float kConfidenceThresh = 0.5;

float *BlobFromImage(cv::Mat &img) {
  float *blob = new float[img.total() * 3];
  int channels = 3;
  int img_h = img.rows;
  int img_w = img.cols;
  for (int c = 0; c < channels; c++) {
    for (int h = 0; h < img_h; h++) {
      for (int w = 0; w < img_w; w++) {
        blob[c * img_w * img_h + h * img_w + w] =
            static_cast<float>(img.at<cv::Vec3b>(h, w)[c]) / 255.0f;
      }
    }
  }
  return blob;
}

inline float GetBoxIoU(const Object &a, const Object &b) {
  cv::Rect2f intersection = a.box & b.box;
  const float i = intersection.area();
  const float u = a.box.area() + b.box.area() - i;
  return (i / u);
}

void SoftNMS(std::vector<Object> *const objects, const float nms_thresh,
             const float confidence_thresh) {
  std::sort(objects->begin(), objects->end(),
            [](const Object &a, const Object &b) {
              return a.confidence > b.confidence;
            });
  std::vector<Object> reserved_objects;
  while (!objects->empty()) {
    const auto obj = objects->front();
    reserved_objects.push_back(obj);
    objects->erase(objects->begin());
    for (auto iter = objects->begin(); iter != objects->end();) {
      const float iou = GetBoxIoU(obj, *iter);
      if (iou > nms_thresh) {
        const float weight = std::exp(-(iou * iou) / 0.5f);
        iter->confidence *= weight;
      }
      if (iter->confidence < confidence_thresh) {
        iter = objects->erase(iter);
      } else {
        ++iter;
      }
    }
  }
  objects->swap(reserved_objects);
}

}  // namespace

bool Yolov6Detector::Init() {
  tensorrt_onnx_.reset(new TensorrtOnnxInference(model_path_));
  if (!tensorrt_onnx_->Init()) {
    std::cout << "Failed to initialize tensorrt onnx inference engine!";
    return -1;
  }
  const auto model_input_dims = tensorrt_onnx_->GetModelInputDims(0);
  model_height_ = model_input_dims.first;
  model_width_ = model_input_dims.second;
  std::cout << "model input height: " << model_height_ << std::endl;
  std::cout << "model input width: " << model_width_ << std::endl;

  return true;
}

bool Yolov6Detector::Detect(cv::Mat &input_image,
                            std::vector<Object> *objects) {
  objects->clear();

  cv::Mat resize_image;
  cv::resize(input_image, resize_image, cv::Size(model_width_, model_height_));
  cv::imwrite("resize.jpg", resize_image);

  float *blob;
  blob = BlobFromImage(resize_image);

  const float width_scale = static_cast<float>(model_width_) / input_image.cols;
  const float height_scale =
      static_cast<float>(model_height_) / input_image.rows;

  std::vector<float *> input_data{blob};
  std::vector<const float *> output_data;

  auto start = std::chrono::steady_clock::now();
  tensorrt_onnx_->Infer(input_data, output_data);
  auto end = std::chrono::steady_clock::now();
  auto cost_time =
      std::chrono::duration_cast<std::chrono::milliseconds>(end - start)
          .count();
  std::cout << "Model inference cost time: " << cost_time << " ms" << std::endl;

  DecodeYolov6Outputs(output_data.at(0), width_scale, height_scale,
                      kConfidenceThresh, objects);
  std::cout << "objects size before nms: " << objects->size() << std::endl;
  SoftNMS(objects, kNmsThresh, kConfidenceThresh);
  std::cout << "objects size after nms: " << objects->size() << std::endl;

  delete[] blob;

  return true;
}

void Yolov6Detector::DecodeYolov6Outputs(const float *const output,
                                         const float width_scale,
                                         const float height_scale,
                                         const float confidence_thresh,
                                         std::vector<Object> *objs) {
  objs->clear();

  float *ptr = const_cast<float *>(output);
  for (int i = 0; i < 8400; ++i, ptr += (kNumClasses + 5)) {
    const float objectness = ptr[4];
    if (objectness >= kObjectnessThresh) {
      const int label =
          std::max_element(ptr + 5, ptr + (kNumClasses + 5)) - (ptr + 5);
      const float confidence = ptr[5 + label] * objectness;
      if (confidence >= confidence_thresh) {
        const float bx = (ptr[0]);
        const float by = (ptr[1]);
        const float bw = ptr[2];
        const float bh = ptr[3];

        Object obj;
        obj.box.x = (bx - bw * 0.5f) / width_scale;
        obj.box.y = (by - bh * 0.5f) / height_scale;
        obj.box.width = bw / width_scale;
        obj.box.height = bh / height_scale;
        obj.label = label;
        obj.confidence = confidence;
        objs->push_back(std::move(obj));
      }
    }
  }  // i loop
}
